// Fill out your copyright notice in the Description page of Project Settings.

#include "WorldHermiteIterator.h"

FWorldHermiteIterator::FWorldHermiteIterator(const std::vector<FWorldRoadGeometry::Ptr>& geoms,
												const FWorldEntityAxis& regionAxis,
												const FWorldEntityAxis& roadAxis,
												const glm::dmat4x4& parentMat)
	: _geoms(geoms)
	, _regionAxis(regionAxis)
	, _roadAxis(roadAxis)
	, _parentMat(parentMat)
	, _geom(!geoms.empty() ? geoms.front()->Variant : nullptr)
	, _geomIdx(0)
	, _intervalIdx(0)
{
	BeginGeometry();
}

bool FWorldHermiteIterator::HasNext() const
{
	return _geom != nullptr;
}

void FWorldHermiteIterator::Next()
{
	if (_geom)
	{
		if (_geom->IsA<FWorldRoadGeometryLine>() ||
			_geom->IsA<FWorldRoadGeometryPoly3>() ||
			_geom->IsA<FWorldRoadGeometryPPoly3>())
		{
			NextGeometry();
		}
		else if (_geom->IsA<FWorldRoadGeometryArc>() ||
				_geom->IsA<FWorldRoadGeometrySpiral>())
		{
			if (_fitEnd < _fitN)
				Polyfit();
			else
				NextGeometry();
		}
		_intervalIdx++;
	}
}

FWorldHermiteIterator::Interval FWorldHermiteIterator::GetInterval()
{
	Interval interval;

	if (_geom)
	{
		// 2D points and vectors on XY plane. (x, y, w)
		glm::dvec3 param, p0, m0, p1, m1;

		if (_geom->IsA<FWorldRoadGeometryLine>())
		{
			p0 = _uv2xy * glm::dvec3(0.0, 0.0, 1.0);
			m0 = _uv2xy * glm::dvec3(_length/2.0, 0.0, 0.0);
			p1 = _uv2xy * glm::dvec3(_length, 0.0, 1.0);
			m1 = _uv2xy * glm::dvec3(_length/2.0, 0.0, 0.0);
			param = glm::dvec3(_intervalIdx, _s, _length);
		}
		else if (_geom->IsA<FWorldRoadGeometryArc>() ||
				_geom->IsA<FWorldRoadGeometrySpiral>())
		{
			glm::dvec2 a = glm::dvec2(_fittedU.x, _fittedV.x);
			glm::dvec2 b = glm::dvec2(_fittedU.y, _fittedV.y);
			glm::dvec2 c = glm::dvec2(_fittedU.z, _fittedV.z);
			glm::dvec2 d = glm::dvec2(_fittedU.w, _fittedV.w);
			p0 = _uv2xy * glm::dvec3(a, 1.0);
			m0 = _uv2xy * glm::dvec3(b, 0.0);
			p1 = _uv2xy * glm::dvec3(a + b + c + d, 1.0);
			m1 = _uv2xy * glm::dvec3(b + 2.0*c + 3.0*d, 0.0);
			param = glm::dvec3(_intervalIdx, _fittedS);
		}
		else if (_geom->IsA<FWorldRoadGeometryPoly3>())
		{
			auto poly3 = _geom->As<FWorldRoadGeometryPoly3>();
			double a = poly3->A;
			double b = poly3->B;
			double c = poly3->C;
			double d = poly3->D;

			// Find U in range [0, range]
			double range = 0.0;
			glm::dvec2 pFrom(0.0, 0.0);
			for (double u = 0.0, s = 0.0; std::abs(s-_length) < 0.1; u += 0.1)
			{
				glm::dvec2 pTo = glm::dvec2(u, a + b*u + c*u*u + d*u*u*u);
				s += glm::length(pTo - pFrom);
				pFrom = pTo;
				range = u;
			}

			glm::dvec2 aa = glm::dvec2(0.0, a);
			glm::dvec2 bb = glm::dvec2(range, b * range);
			glm::dvec2 cc = glm::dvec2(0.0, c * range * range);
			glm::dvec2 dd = glm::dvec2(0.0, d * range * range * range);

			_roadAxis.ToProjPoly(aa, bb, cc, dd, _regionAxis);

			p0 = _uv2xy * glm::dvec3(aa, 1.0);
			m0 = _uv2xy * glm::dvec3(bb, 0.0);
			p1 = _uv2xy * glm::dvec3(aa + bb + cc + dd, 1.0);
			m1 = _uv2xy * glm::dvec3(bb + 2.0*cc + 3.0*dd, 0.0);
			param = glm::dvec3(_intervalIdx, _s, _length);
		}
		else if (_geom->IsA<FWorldRoadGeometryPPoly3>())
		{
			auto ppoly3 = _geom->As<FWorldRoadGeometryPPoly3>();
			glm::dvec2 a(ppoly3->Au, ppoly3->Av);
			glm::dvec2 b(ppoly3->Bu, ppoly3->Bv);
			glm::dvec2 c(ppoly3->Cu, ppoly3->Cv);
			glm::dvec2 d(ppoly3->Du, ppoly3->Dv);

			if (ppoly3->Prange == FWorldPPolyParameterRange::ArcLength)
			{
				b *= _length;
				c *= _length * _length;
				d *= _length * _length * _length;
			}

			_roadAxis.ToProjPoly(a, b, c, d, _regionAxis);

			p0 = _uv2xy * glm::dvec3(a, 1.0);
			m0 = _uv2xy * glm::dvec3(b, 0.0);
			p1 = _uv2xy * glm::dvec3(a + b + c + d, 1.0);
			m1 = _uv2xy * glm::dvec3(b + 2.0*c + 3.0*d, 0.0);
			param = glm::dvec3(_intervalIdx, _s, _length);
		}

		// Put 2D points in 3D space with z=0.0
		interval.Param = param;
		interval.P0 = glm::dvec3(_parentMat * glm::dvec4(p0.x, p0.y, 0.0, 1.0));
		interval.M0 = glm::dvec3(_parentMat * glm::dvec4(m0.x, m0.y, 0.0, 0.0));
		interval.P1 = glm::dvec3(_parentMat * glm::dvec4(p1.x, p1.y, 0.0, 1.0));
		interval.M1 = glm::dvec3(_parentMat * glm::dvec4(m1.x, m1.y, 0.0, 0.0));

		// Bounding box of control points..
		interval.BBox |= interval.P0;
		interval.BBox |= interval.P1;
	}

	return interval;
}

void FWorldHermiteIterator::NextGeometry()
{
	if (_geom)
	{
		_geomIdx++;
		_geom = _geomIdx < _geoms.size() ? _geoms[_geomIdx]->Variant : nullptr;
		BeginGeometry();
	}
}

void FWorldHermiteIterator::BeginGeometry()
{
	if (_geom)
	{
		// Change this to control the precision level of arc and spiral fitting algorithm.
		constexpr double arcSampleRes = 0.02;
		constexpr double arcMaxSampleRes = 10.0;
		constexpr double arcMinStep = 16;

		FWorldRoadGeometry::Ptr geom = _geoms[_geomIdx];

		_fitU.clear();
		_fitV.clear();
		_fitT.clear();
		_fitS.clear();
		_fitN = _fitBegin = _fitEnd = 0;

		_pointRaw = glm::dvec3(geom->X, geom->Y, 0.0);
		_pointOnXY = _regionAxis.ToProj(_pointRaw);
		_hdg = geom->Hdg;
		_s = geom->S;
		_length = geom->Length;

		_uv2xy = glm::dmat3x3(
			std::cos(_hdg),	std::sin(_hdg),	0.0,
			-std::sin(_hdg),std::cos(_hdg),	0.0,
			_pointOnXY.x,	_pointOnXY.y,	1.0
		);

		if (_geom->IsA<FWorldRoadGeometryArc>())
		{
			auto arc = _geom->As<FWorldRoadGeometryArc>();
			double curvature = arc->Curvature;

			// Adaptive step resolution (depends on curvature).
			double arcRes = arcSampleRes / std::abs(curvature);
			arcRes = glm::min(arcMaxSampleRes, arcRes);

			// Sample linearly along the run length.
			double arcStep = glm::ceil(_length / arcRes);
			arcStep = glm::max(arcMinStep, arcStep);

			_arcIt.Reset(curvature, _length, _length / arcStep + DBL_EPSILON);
			for (; _arcIt.HasNext(); _arcIt.Next())
			{
				glm::dvec3 p = _arcIt.Sample();

				_fitU.push_back(p.x);
				_fitV.push_back(p.y);
				_fitT.push_back(0.0);
				_fitS.push_back(_arcIt.S());
				_fitN++;
			}
			Polyfit();
		}
		else if (_geom->IsA<FWorldRoadGeometrySpiral>())
		{
			auto spiral = _geom->As<FWorldRoadGeometrySpiral>();
			double curvStart = spiral->CurvStart;
			double curvEnd = spiral->CurvEnd;

			// Adaptive step resolution (depends on curvature).
			double spiralRes = arcSampleRes * 0.5 / std::abs(curvStart - curvEnd);
			spiralRes = glm::min(arcMaxSampleRes, spiralRes);

			// Sample linearly along the run length.
			double spiralStep = glm::ceil(_length / spiralRes);
			spiralStep = glm::max(arcMinStep, spiralStep);

			_spiralIt.Reset(curvStart, curvEnd, _length, _length / spiralStep + DBL_EPSILON);
			for (; _spiralIt.HasNext(); _spiralIt.Next())
			{
				glm::dvec3 p = _spiralIt.Sample();

				_fitU.push_back(p.x);
				_fitV.push_back(p.y);
				_fitT.push_back(0.0);
				_fitS.push_back(_spiralIt.S());
				_fitN++;
			}
			Polyfit();
		}
	}
}

void FWorldHermiteIterator::Polyfit()
{
	// Allowed distance between fitted point and point on arc/spiral.
	constexpr double MaxDiff = 0.001;

	// Resume fitting from the last point.
	_fitBegin = std::max(_fitBegin, _fitEnd-1);

	// Hungry fit and binary search for fitness.
	_fitEnd = _fitN;

	// Polyfit if we have 4 points at least.
	while (_fitEnd - _fitBegin >= 4)
	{
		// t in range[0, 1], linearly along the road.
		for (int i = _fitBegin; i < _fitEnd; i++)
		{
			_fitT[i] = (i-_fitBegin) / (double)(_fitEnd-_fitBegin-1);
		}

		// Fit points in range [fitBegin, fitEnd)
		double polyU[4], polyV[4];
		int okX = FWorldMath::Polyfit(_fitT.data()+_fitBegin, _fitU.data()+_fitBegin, _fitEnd-_fitBegin, 3, polyU);
		int okY = FWorldMath::Polyfit(_fitT.data()+_fitBegin, _fitV.data()+_fitBegin, _fitEnd-_fitBegin, 3, polyV);

		if (okX == 0 && okY == 0)
		{
			bool fitOk = true;

			glm::dvec2 a(polyU[0], polyV[0]);
			glm::dvec2 b(polyU[1], polyV[1]);
			glm::dvec2 c(polyU[2], polyV[2]);
			glm::dvec2 d(polyU[3], polyV[3]);

			// Check
			for (int i = _fitBegin; i < _fitEnd; i++)
			{
				double t = _fitT[i];
				double t2 = t * t;
				double t3 = t * t2;

				glm::dvec2 pfit = a + b*t + c*t2 + d*t3;
				glm::dvec2 psrc = glm::dvec2(_fitU[i], _fitV[i]);
				if (glm::distance(pfit, psrc) > MaxDiff)
				{
					fitOk = false;
					break;
				}
			}

			// Meet the distance criteria.
			if (fitOk)
			{
				double s = _s + _fitS[_fitBegin];
				double length = _fitS[_fitEnd-1] - _fitS[_fitBegin];
				_fittedU = glm::dvec4(a.x, b.x, c.x, d.x);
				_fittedV = glm::dvec4(a.y, b.y, c.y, d.y);
				_fittedS = glm::dvec2(s, length);
				break;
			}
		}

		// Not satisfied. Binary search the first half. Remember that we
		// have at least 4 points in the range.
		_fitEnd = std::max(_fitEnd-(_fitEnd-_fitBegin)/2, _fitBegin+4);
	}

	// Fit the remaining points. Above while-loop fits >=4 points. We fit
	// from the last point (fitEnd-1) so we should have at least 2 points.
	// i.e. 2 or 3 points remaining.
	if (_fitEnd - _fitBegin > 1 && _fitEnd - _fitBegin < 4)
	{
		int degree = _fitEnd - _fitBegin - 1;

		// t in range[0, 1], linearly along the remaining.
		for (int i = _fitBegin; i < _fitEnd; i++)
		{
			_fitT[i] = (i-_fitBegin) / (double)(_fitEnd-_fitBegin-1);
		}

		// Fit points in range [fitBegin, fitEnd)
		double polyU[4], polyV[4];
		int okX = FWorldMath::Polyfit(_fitT.data()+_fitBegin, _fitU.data()+_fitBegin, _fitEnd-_fitBegin, degree, polyU);
		int okY = FWorldMath::Polyfit(_fitT.data()+_fitBegin, _fitV.data()+_fitBegin, _fitEnd-_fitBegin, degree, polyV);

		if (degree < 3) polyU[3] = polyV[3] = 0.0;
		if (degree < 2) polyU[2] = polyV[2] = 0.0;

		double s = _s + _fitS[_fitBegin];
		double length = _fitS[_fitEnd-1] - _fitS[_fitBegin];

		if (okX == 0 && okY == 0)
		{
			glm::dvec2 a(polyU[0], polyV[0]);
			glm::dvec2 b(polyU[1], polyV[1]);
			glm::dvec2 c(polyU[2], polyV[2]);
			glm::dvec2 d(polyU[3], polyV[3]);

			_fittedU = glm::dvec4(a.x, b.x, c.x, d.x);
			_fittedV = glm::dvec4(a.y, b.y, c.y, d.y);
			_fittedS = glm::dvec2(s, length);
		}
		else
		{
			// Not likely to happen, straight line in case of failure.
			glm::dvec2 p0(_fitU[_fitBegin], _fitV[_fitBegin]);
			glm::dvec2 p1(_fitU[_fitEnd-1], _fitV[_fitEnd-1]);
			glm::dvec2 m0 = p1 - p0;
			glm::dvec2 m1 = m0;

			glm::dvec2 a = p0;
			glm::dvec2 b = m0;
			glm::dvec2 c = -3.0 * p0 - 2.0 * m0 + 3.0 * p1 - m1;
			glm::dvec2 d = 2.0 * p0 + m0 - 2.0 * p1 + m1;

			_fittedU = glm::dvec4(a.x, b.x, c.x, d.x);
			_fittedV = glm::dvec4(a.y, b.y, c.y, d.y);
			_fittedS = glm::dvec2(s, length);
		}
	}
}

